# -*- coding: utf-8 -*-
"""
@Time: 2020/12/14 14:16
@Author: 鹄望潇湘
@File: PascalDataSet.py
@Desc: 本文件包含所有对PascalVOC数据集的相关操作，其中PascalLabelImageDataSet类用来将标注好的图片的每个像素转换成对应的类别
编号后生成DataSet， ImageScaleCrop类用于将图片周围进行填充使其长，宽像素变为某个值的整数倍，这样做的目的实际上是为了使得反卷积
前后图片的尺寸一致，同时使网络能够适应不同尺寸的图片。
"""

from torchvision.transforms import Compose
from torchvision import transforms
from tools_zsj.ZsjDataSet import MultiDataSet, SingleImageDataSet, SingleDataSet

import os
import pickle
import numpy as np
from PIL.Image import Image
import torch


class PascalLabelImageDataSet(SingleImageDataSet):
    """
    为Pascal VOC数据集设计的已经标注图片转类别的DataSet类。

    指定Pascal VOC数据集对应的类别字典，然后添加SegmentationClass文件夹下的所有图片即可完成配置。

    调用__getitem__函数时，会将对应的标注图片的每一个像素转换成对应的类别，最终返回一个与图片尺寸相同的矩阵。矩阵中每一个元素都是
    图片内容中的每个像素点对应的类别。

    :param color_object_json_path: color_object_dict.pickle文件的存放位置，该文件可使用create_color_objet_dict.py生成
    """

    def __init__(self, color_object_json_path: str):
        super(PascalLabelImageDataSet, self).__init__()
        self.color_object_dict_path = color_object_json_path
        self.color_object_dict = dict()
        self.__read_color_object_dict()
        # self.crop_transformer = ImageScaleCrop(32)

    def __getitem__(self, item):
        image_data = self.read_image_data(self.data[item])
        image_data = self.transformers(image_data)

        image_data = np.array(image_data)
        result = np.apply_along_axis(self.__get_class_from_dict, 2, image_data)
        result = torch.from_numpy(result.astype(np.int64))
        return result

    def __get_class_from_dict(self, vec):
        vec = vec.astype(np.uint8)
        label = vec.tostring()
        classes = self.color_object_dict[label]
        return classes

    def __read_color_object_dict(self):
        """
        从color_object_dict.pickle文件中读取内容，保存成字典

        :return: None
        """
        with open(self.color_object_dict_path, 'rb') as f:
            self.color_object_dict = pickle.load(f)
        pass


class NormalizedLabelDataSet(SingleDataSet):
    def __getitem__(self, item):
        return self.__load_pickle(self.data[item])[0]

    def __load_pickle(self, path: str) -> torch.Tensor:
        with open(path, "rb") as f:
            result = pickle.load(f)
            return result

    def push_labeled_image(self, path: str):
        self.add_data(path)


class ImageScaleCrop(object):
    """
    将图片的宽度和高度填充到指定值(scale)的倍数，保证卷积与反卷积后的尺寸大小是一致的

    例如：当scale的值为32时， 510*253的图片将被填充至512*256

    :param scale 指定的倍数值

    """
    def __init__(self, scale: int = 32):
        self.scale = scale

    def __call__(self, pic: Image):
        """

        :param pic:
        :type pic Image
        :return:
        """
        if not isinstance(pic, Image):
            return pic
        width = pic.width
        height = pic.height

        target_width = width
        target_height = height
        if width % self.scale != 0:
            target_width = (width // self.scale + 1) * self.scale
        if height % self.scale != 0:
            target_height = (height // self.scale + 1) * self.scale

        transform = transforms.CenterCrop((target_height, target_width))
        result = transform(pic)
        return result


class PascalDataSetGenerator(object):
    """
    Pascal VOC数据集的训练集与验证集生成器，构造函数参数为:

    color_object_dict_path: Pascal VOC类别字典的文件路径(color_object_dict.pickle文件路径)

    :param color_object_dict_path: Pascal VOC类别字典的文件路径（color_object_dict.pickle文件路径）
    :return: None
    """

    def __init__(self, color_object_dict_path: str):
        self.color_object_dict_path = color_object_dict_path

    def create_dataset_by_txt(self, txt_file_path: str, original_image_path: str,
                              labeled_image_path: str, crop_size: tuple = None) -> MultiDataSet:
        """
        生成Pascal VOC的图片生成神经网络训练需要使用的数据集，返回类型为MultiDataSet,
        其中包含了原图的RGB数据，以及该图片对应标注的图片的类别数据矩阵

        :param crop_size:
        :param txt_file_path: 要生成的数据集的图片文件名列表文件的存放路径
        :param original_image_path: 真实图片存放的文件夹路径
        :param labeled_image_path: 已经标注好的图片的文件夹路径　
        :return: 一个MultiDataSet对象，包含了原图以及对应的每个像素的类别矩阵
        :rtype: MultiDataSet
        """
        result = MultiDataSet()
        image_dataset, value_dataset = self.__prepare_pair_dataset(crop_size)

        with open(txt_file_path, "r") as f:
            for line in f:
                line = line.strip('\n')
                ori_image_path = os.path.join(original_image_path, line + ".jpg")
                lab_image_path = os.path.join(labeled_image_path, line + ".png")
                if os.path.isfile(ori_image_path) and os.path.isfile(lab_image_path):
                    image_dataset.push_single_image(ori_image_path)
                    value_dataset.push_single_image(lab_image_path)

        result.append(image_dataset)
        result.append(value_dataset)
        return result

    def create_normalized_dataset(self, txt_file_path: str, original_image_path: str,
                                  labeled_image_path: str) -> MultiDataSet:
        result = MultiDataSet()
        image_transformer = Compose([transforms.ToTensor(),
                                     transforms.Normalize([0.5, 0.5, 0.5],
                                                          [0.5, 0.5, 0.5])])
        image_dataset = SingleImageDataSet()
        value_dataset = NormalizedLabelDataSet()
        image_dataset.set_image_transformer(image_transformer)

        with open(txt_file_path, "r") as f:
            for line in f:
                line = line.strip('\n')
                ori_image_path = os.path.join(original_image_path, line + ".jpg")
                lab_image_path = os.path.join(labeled_image_path, line + ".pickle")
                if os.path.isfile(ori_image_path) and os.path.isfile(lab_image_path):
                    image_dataset.push_single_image(ori_image_path)
                    value_dataset.push_labeled_image(lab_image_path)
        result.append(image_dataset)
        result.append(value_dataset)
        return result

    def create_dataset_by_image(self, real_image_path: str, label_image_path: str) -> MultiDataSet:
        """

        :param real_image_path: 真实图片存放路径
        :param label_image_path: 对应的标注图片存放路径
        :return:
        """
        result = MultiDataSet()
        image_dataset, value_dataset = self.__prepare_pair_dataset()

        image_dataset.push_single_image(real_image_path)
        value_dataset.push_single_image(label_image_path)

        result.append(image_dataset)
        result.append(value_dataset)
        return result

    def __prepare_pair_dataset(self, crop_size: tuple = None) -> (SingleImageDataSet, PascalLabelImageDataSet):
        image_dataset = SingleImageDataSet()
        value_dataset = PascalLabelImageDataSet(color_object_json_path=self.color_object_dict_path)

        # 设置transformer, 将图像数据与转换成Tensor
        image_transformer = None
        label_transformer = None
        if crop_size is None:
            image_transformer = Compose([ImageScaleCrop(32),
                                        transforms.ToTensor(),
                                        transforms.Normalize([0.5, 0.5, 0.5],
                                                            [0.5, 0.5, 0.5])
                                         ])
            label_transformer = Compose([ImageScaleCrop(32)])
        else:
            image_transformer = Compose([transforms.CenterCrop(crop_size),
                                         transforms.ToTensor(),
                                         transforms.Normalize([0.5, 0.5, 0.5],
                                                              [0.5, 0.5, 0.5])
                                         ])
            label_transformer = Compose([transforms.CenterCrop(crop_size)])
        image_dataset.set_image_transformer(image_transformer)
        value_dataset.set_image_transformer(label_transformer)

        return image_dataset, value_dataset
